import os

import cv2
import numpy as np
from torch.utils.data import Dataset
np.set_printoptions(threshold=np.inf)


class BaseDataset(Dataset):

    def __init__(self, path, mode='train'):
        self.data_path = path
        self.img_paths = []
        self.annotation = []
        self.sizes = []
        self.cur_indexes = []
        self.mode = mode
        self.data_pipe = []
        self.addition_op = None
        self.load_data()

    def load_data(self):
        raise Exception('BaseDataset does not implement load_data function')

    def __len__(self):
        return len(self.cur_indexes)

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        img = cv2.imread(os.path.join(self.data_path, img_path))
        # To rgb
        r = img[:, :, 2]
        img[:, :, 2] = img[:, :, 0]
        img[:, :, 0] = r

        annotation = self.annotation[idx]

        size = self.sizes[idx]
        if self.addition_op:
            self.addition_op(idx, img)

        if img.dtype != 'float32':
            img = img.astype(np.float32)
        for preprocess_op in self.data_pipe:
            img, annotation = preprocess_op.process(img, annotation)
        img = img.transpose(2, 0, 1)
        return img, annotation, size, img_path

    def set_data_pipe(self, data_pipe):
        self.data_pipe = data_pipe

